#!/bin/sh
# @Author duangr
# @Date   2013-12-16
# @Blog   http://my.oschina.net/duangr
# ---------------------------------------
# enable environment config 
source /etc/profile
source ~/.bash_profile

# ---------------------------------------
# find BIN_HOME
BIN_HOME=`dirname "$0"`
BIN_HOME=`cd "$BIN_HOME"; pwd`
SHELL_NAME=$(basename $0)

# ---------------------------------------
author="@duangr"
version="1.0"

# declare return code
RTN_OK=0
RTN_WARNING=1
RTN_CRITICAL=2
RTN_UNKNOWN=3

# declare params
p_warn=
p_critical=
p_port=
p_state=

# declare info and result
INFO=
RESULT=

# ---------------------------------------
# function declare area
_f_validate(){
  k=$1
  v=$2
  if [ "$v" = "" ]; then
    echo "Invalid ${k} Param: NULL POINT"
    exit 1
  fi
  if [ "${v:0:1}" = "-" ]; then
    echo "Invalid ${k} Param: ${v}"
    exit 1
  fi
}
_f_read_params(){
  pnum=$#
  i=1
  while [ $i -le $pnum ];
  do
    v=${!i}

    if [ "$v" = "-h" ] || [ "$v" = "--help" ]; then
      _f_help
      exit 0
    fi
    if [ "$v" = "-V" ] || [ "$v" = "--version" ]; then
      _f_version
      exit 0
    fi

    if [ "$v" = "-w" ]; then
      nI=`expr $i + 1`
      nV=${!nI}
      _f_validate "${v}" "${nV}"
      i=`expr $i + 1`
      p_warn=${nV}
    fi

    if [ "$v" = "-c" ]; then
      nI=`expr $i + 1`
      nV=${!nI}
      _f_validate "${v}" "${nV}"
      i=`expr $i + 1`
      p_critical=${nV}
    fi

    if [ "$v" = "-p" ]; then
      nI=`expr $i + 1`
      nV=${!nI}
      _f_validate "${v}" "${nV}"
      i=`expr $i + 1`
      p_port=${nV}
    fi


    if [ "$v" = "-s" ]; then
      nI=`expr $i + 1`
      nV=${!nI}
      _f_validate "${v}" "${nV}"
      i=`expr $i + 1`
      p_state=${nV}
    fi

    # increase
    i=`expr $i + 1`
  done;
}

_f_check(){
  CMD="netstat -an"
  INFO="netstat"
  if [ -n "${p_state}" ]; then
    CMD="${CMD} |grep ${p_state}"
    INFO="${INFO} State: ${p_state}"
  fi
  if [ -n "${p_port}" ]; then
    CMD="${CMD} |grep \":${p_port} \""
    INFO="${INFO} Port: ${p_port}"
  fi
  CMD="${CMD} |wc -l"
  RESULT=`echo $CMD | /bin/sh`
  INFO="${INFO} Num: ${RESULT}"
}

_f_threshold(){
  # check threshold
  # validate result is a number
  if [ `echo ${RESULT} |grep -P "^-?\d+$" |wc -l` != 1 ]; then
    echo "UNKNOWN STATE"
    exit ${RTN_UNKNOWN}
  fi

  # warning
  t_w_max=0
  t_w_min=0
  if [[ -n $p_warn ]]; then
    t_val=$p_warn
    colon_index=`expr index "$t_val" :`
    if [[ $colon_index -eq 0 ]]; then
      t_w_max=$t_val
    else
      t_w_min=${t_val:0:`expr $colon_index - 1`}
      t_w_max=${t_val:colon_index}
    fi
  fi

  # critical
  t_c_max=0
  t_c_min=0
  if [[ -n $p_critical ]]; then
    t_val=$p_critical
    colon_index=`expr index "$t_val" :`
    if [[ $colon_index -eq 0 ]]; then
      t_c_max=$t_val
    else
      t_c_min=${t_val:0:`expr $colon_index - 1`}
      t_c_max=${t_val:colon_index}
    fi
  fi

  # check start
  if [[ ${RESULT} -lt ${t_c_min} ]] || [[ ${RESULT} -gt ${t_c_max} ]]; then
    echo "NETSTAT CRITICAL: ${INFO}" 
    exit ${RTN_CRITICAL}
  elif [[ ${RESULT} -lt ${t_w_min} ]] || [[ ${RESULT} -gt ${t_w_max} ]]; then
    echo "NETSTAT WARNING: ${INFO}"
    exit ${RTN_WARNING}
  else
    echo "NETSTAT OK: ${INFO}"
    exit ${RTN_OK}
  fi
}
_f_version(){
  echo "${SHELL_NAME} ${version} ${author}"
}
_f_help(){
  _f_version
  echo ""
  echo "Checks all network connects and generates WARNING or CRITICAL states if the specified"
  echo "metric is outside the required threshold ranges. The metric defaults to number"
  echo "of connects.  Search filters can be applied to limit the connects to check."
  echo ""
  echo "Usage: ${SHELL_NAME} -w limit -c limit [-p Port] [-s State]"
  echo ""
  echo "Options:"
  echo " -h, --help"
  echo "    Print detailed help screen"
  echo " -V, --version"
  echo "    Print version information"
  echo " -w"
  echo "    Generate warning state if metric is outside this range"
  echo " -c"
  echo "    Generate critical state if metric is outside this range"
  echo " -p"
  echo "    Port Num for filtering"
  echo " -s"
  echo "    State for filtering (e.g. ESTABLISHED LISTEN TIME_WAIT)"
  echo ""
  echo "Examples:"
  echo " ${SHELL_NAME} -w 2 -c 4 -p 22 -s ESTABLISHED"
  echo "   Critical if > 4 ESTABLISHED connects with port 22"
  echo "   Warning  if > 2 ESTABLISHED connects with port 22"
  echo ""
  echo " ${SHELL_NAME} -w 2:10 -c :20 -p 80 -s ESTABLISHED"
  echo "   Critical if > 20 ESTABLISHED connects with port 80"
  echo "   Warning  if < 2 or > 10 ESTABLISHED connects with port 80"
  echo ""
  echo " ${SHELL_NAME} -c 1:1 -p 80 -s LISTEN"
  echo "   Critical if not one LISTEN connects with port 80"
  echo ""
}

# ----------------------------
# main
_f_read_params $*
_f_check
_f_threshold
